from __future__ import division, print_function
from svgpathtools import (
    Path,
    Line,
    CubicBezier,
    smoothed_path,
)
import svgwrite
import pandas as pd
import numpy as np
from consts import *


class DrawEngine:
    def __init__(self) -> None:
        df = pd.read_excel("data.xlsx", sheet_name="product")
        df["val"] = df.val / RATIO
        self.df = df
        self.node_df = pd.read_excel("data.xlsx", sheet_name="node")
        self.line_df = pd.read_excel("data.xlsx", sheet_name="line")
        self.dwg = svgwrite.Drawing(
            "map.svg", size=CAVAS_SIZE, debug=True, style="background-color:white"
        )
        self.calc_node_df = pd.DataFrame()
        self.used_df = pd.DataFrame()
        self.end_node_df = self.node_df[self.node_df.cat == "终端"].reset_index(drop=True)
        self.end_y = START_POS[1] + 50

    def run_first_layer(self, df, cat, cluster_cat):
        dwg = self.dwg
        df = df.copy().reset_index(drop=True)
        total = df.val.sum()
        node_df = self.node_df
        node_df = node_df[node_df.cat == cat].copy().reset_index(drop=True)
        x, y = node_df[node_df.name == cat][["x", "y"]].values[0]
        df["height"] = np.clip(df.val, a_max=None, a_min=FONT_SIZE)
        height_arr = df.height.values
        df["margin"] = np.arange(len(height_arr)) * MARGIN
        height_arr = np.concatenate([[0], height_arr])
        height_arr = np.cumsum(height_arr)
        total_height = height_arr[-1] + df.margin.values[-1]
        off_height = total_height - total
        height_arr = height_arr
        height_arr = height_arr[:-1]
        df["cum_val"] = np.cumsum(np.concatenate([[0], df.val.values[:-1]]))
        df["cum_height"] = height_arr + df.margin
        df["y1"] = df.cum_height + df.height / 2.0 + y
        df["x1"] = x + df.margin.values[::-1] + 10
        df["x2"] = 400 + x
        df["y2"] = df.cum_val + df.val / 2.0 + off_height + y
        all_paths = []
        all_vals = []
        # df.to_clipboard()
        end_y = 0.0
        for i, row in df.iterrows():
            seg0 = Line(complex(x, row.y1), complex(row.x1, row.y1))
            # seg1 = Line(complex(row.x1, row.y1), complex(row.x2, row.y2))
            seg1 = CubicBezier(
                complex(row.x1, row.y1),
                complex(row.x1 / 2.0 + row.x2 / 2.0, row.y1),
                complex(row.x1 / 2.0 + row.x2 / 2.0, row.y2),
                complex(row.x2, row.y2),
            )
            seg2 = Line(complex(row.x2, row.y2), complex(row.x2 + 100, row.y2))
            path = Path(seg0, seg1, seg2)
            spath = smoothed_path(path, tightness=0.5)
            all_paths.append(spath)

            all_vals.append(
                dwg.text(
                    str(np.round(row.val, 1)),
                    insert=(x + 10, row.y1 + 6),
                    fill="black",
                    style="font-weight:bolder",
                )
            )
            all_vals.append(
                dwg.text(
                    row["name"],
                    insert=(x - 72, row.y1 + 6),
                    fill="black",
                    style="font-weight:bolder",
                )
            )
            if i == df.shape[0] - 1:
                end_y = row.y2 + row.val / 2.0
        end_x, end_y = df.x2.values[0], end_y - total / 2.0

        end_path = Line(complex(end_x, end_y), complex(end_x + 100, end_y))

        all_paths.append(Path(end_path))

        df = df.merge(self.node_df, on=["name", "cat"])
        colors = list(df.color.values) + list(
            node_df[node_df.name == cluster_cat].color.values
        )
        names = list(df["name"].values) + [""]

        calc_node_arr = [
            {
                "name": cluster_cat,
                "bottom": end_path.end + complex(0, total / 2.0),
                "top": end_path.end + complex(0, -total / 2.0),
            }
        ]
        self.calc_node_df = pd.concat([self.calc_node_df, pd.DataFrame(calc_node_arr)])

        width_arr = list(df.val.values) + [total]
        power_df = pd.DataFrame()
        power_df["path"] = [p.d() for p in all_paths]
        power_df["name"] = names
        power_df["color"] = colors
        power_df["width"] = width_arr

        for i, row in power_df.iterrows():
            dwg.add(
                dwg.path(
                    d=row.path, stroke=row.color, fill="none", stroke_width=row.width
                )
            )
        for val in all_vals:
            dwg.add(val)

        cat_y = y + total_height if cat == "外来电" else y + total_height / 2.0 - 9
        dwg.add(
            dwg.text(
                cat,
                insert=(x - 170, cat_y),
                fill="black",
                style="font-weight:bolder",
            )
        )
        end_path = calc_node_arr[0]
        print(end_path)
        dwg.add(
            dwg.text(
                str(np.round(total, 1)),
                insert=(
                    end_path["top"].real - 75,
                    end_path["top"].imag + total / 2.0 + FONT_SIZE / 2.0,
                ),
                fill="black",
                style="font-weight:bolder",
            )
        )
        p0 = complex(x - 90, y)
        p1 = complex(x - 100, p0.imag)
        p2 = complex(p1.real, p1.imag + total_height)
        p3 = complex(p0.real, p2.imag)
        seg0 = Line(p0, p1)
        seg1 = Line(p1, p2)
        seg2 = Line(p2, p3)
        left_path = Path(seg0, seg1, seg2)
        dwg.add(
            dwg.path(
                d=left_path.d(),
                stroke="black",
                fill="none",
                stroke_width=1,
                stroke_linecap="round",
            )
        )

    def run_second_layer(self, to):
        dwg = self.dwg
        line_df = (
            self.line_df[(self.line_df.layer == 2) & (self.line_df.to == to)]
            .copy()
            .reset_index(drop=True)
        )
        line_df.val = line_df.val / RATIO
        calc_node_df = self.calc_node_df
        line_df = line_df.merge(calc_node_df, left_on="from", right_on="name")
        node_df = self.node_df.copy()
        total_width = line_df.val.sum()
        to_row = node_df[node_df.name == to].iloc[0]
        last_line = line_df.loc[line_df.shape[0] - 1]
        to_x = to_row.x
        to_y2 = to_y = (
            last_line.top.imag - 30 - total_width + last_line.val
            if last_line.from_direction == "top"
            else last_line.bottom.imag + 30 - total_width + last_line.val
        )
        if line_df.shape[0] == 1 and last_line.from_direction == "bottom":
            to_y = to_y2 = last_line.bottom.imag - last_line.val + 30

        data = []

        for i, row in line_df.iterrows():
            start = row.top if row.from_direction == "top" else row.bottom
            ratio = 1 if row.from_direction == "top" else -1
            p0 = start + complex(0, ratio * row.val / 2.0)
            p1 = start + complex(15, ratio * row.val / 2.0)
            p4 = complex(to_x, to_y + row.val / 2.0)

            p2 = complex(p1.real / 2.0 + p4.real / 2.0, p1.imag)
            p3 = complex(p1.real / 2.0 + p4.real / 2.0, p2.imag)
            p5 = p4 + complex(20, 0)
            seg0 = Line(p0, p1)
            seg1 = CubicBezier(p1, p2, p3, p4)
            seg2 = Line(p4, p5)
            path = Path(seg0, seg1, seg2)
            to_y = to_y + row.val
            if row.from_direction == "top":
                used_df = pd.DataFrame(
                    [
                        {
                            "name": row["name"],
                            "top": p0 - complex(0, row.val / 2.0),
                            "bottom": p0 + complex(0, row.val / 2.0),
                        }
                    ]
                )
                self.used_df = pd.concat([self.used_df, used_df], ignore_index=True)
            # spath = smoothed_path(path, tightness=0.5)
            data.append(
                {
                    "path": path.d(),
                    "width": row.val,
                    "color": to_row.color,
                    "name": row["name"],
                }
            )

        draw_df = pd.DataFrame(data)
        for _, row in draw_df.iterrows():
            dwg.add(
                dwg.path(
                    d=row.path,
                    stroke=row.color,
                    fill="none",
                    stroke_width=row.width,
                    style="font-weight:bolder;",
                )
            )
        p6 = complex(to_x + 10, to_y2 + total_width / 2.0)
        end_path = Path(Line(p6, p6 + complex(180, 0)))
        dwg.add(
            dwg.path(
                d=end_path.d(), stroke="green", fill="none", stroke_width=total_width
            )
        )

        if len(to) == 4:
            dwg.add(
                dwg.text(
                    to[0] + " " + to[1],
                    insert=(to_x + 90 - 36, to_y2 + total_width / 2.0 - 24),
                    fill="white",
                    font_size=36,
                    style="font-weight:bolder;",
                )
            )
            dwg.add(
                dwg.text(
                    to[2] + " " + to[3],
                    insert=(to_x + 90 - 36, to_y2 + total_width / 2.0 + 24),
                    fill="white",
                    font_size=36,
                    style="font-weight:bolder;",
                )
            )
        else:
            dwg.add(
                dwg.text(
                    to[0] + " " + to[1],
                    insert=(to_x + 90 - 27, to_y2 - ratio * total_width / 2.0 + 9),
                    fill="white",
                    font_size=36,
                    style="font-weight:bolder;",
                )
            )
        calc_node_arr = [
            {
                "name": to,
                "bottom": end_path.end + complex(0, total_width / 2.0),
                "top": end_path.end + complex(0, -total_width / 2.0),
            }
        ]
        self.calc_node_df = pd.concat([self.calc_node_df, pd.DataFrame(calc_node_arr)])

    def run_third_layer(self, to):
        dwg = self.dwg
        line_df = (
            self.line_df[(self.line_df.layer == 3) & (self.line_df.to == to)]
            .copy()
            .reset_index(drop=True)
        )
        line_df.val = line_df.val / RATIO
        calc_node_df = self.calc_node_df
        # print(calc_node_df)
        line_df = line_df.merge(calc_node_df, left_on="from", right_on="name")
        node_df = self.node_df.copy()
        total_width = line_df.val.sum()
        to_row = node_df[node_df.name == to].iloc[0]
        last_line = line_df.loc[line_df.shape[0] - 1]
        to_x = to_row.x
        to_y2 = to_y = (
            last_line.top.imag - 30 - total_width + last_line.val
            if last_line.from_direction == "top"
            else last_line.bottom.imag + 30 - total_width + last_line.val
        )
        if line_df.shape[0] == 1 and last_line.from_direction == "bottom":
            to_y2 = to_y = last_line.bottom.imag - last_line.val + 30
        # print(line_df)
        # print(f"top:{last_line.top.imag} - 30 - {total_width} + {last_line.val}")
        # print(f"bottom:{last_line.bottom.imag} - 30 - {total_width} + {last_line.val}")
        # print(f"to_x: {to_x},to_y: {to_y}, direction:{last_line.from_direction}")
        data = []

        for i, row in line_df.iterrows():
            start = row.top if row.from_direction == "top" else row.bottom
            ratio = 1 if row.from_direction == "top" else -1
            p0 = start + complex(0, ratio * row.val / 2.0)
            p1 = start + complex(15, ratio * row.val / 2.0)
            p4 = complex(to_x, to_y + row.val / 2.0)

            # print(f"p1:{p1}, p4:{p4}")
            p2 = complex(p1.real / 2.0 + p4.real / 2.0, p1.imag)
            p3 = complex(p1.real / 2.0 + p4.real / 2.0, p2.imag)
            p5 = p4 + complex(20, 0)
            seg0 = Line(p0, p1)
            seg1 = CubicBezier(p1, p2, p3, p4)
            seg2 = Line(p4, p5)
            path = Path(seg0, seg1, seg2)
            path = smoothed_path(path, tightness=0.05)
            to_y = to_y + row.val
            used_df = pd.DataFrame(
                [
                    {
                        "name": row["name"],
                        "top": p0 - complex(0, row.val / 2.0),
                        "bottom": p0 + complex(0, row.val / 2.0),
                    }
                ]
            )
            self.used_df = pd.concat([self.used_df, used_df], ignore_index=True)
            # spath = smoothed_path(path, tightness=0.5)
            data.append({"path": path.d(), "width": row.val, "color": to_row.color})
        draw_df = pd.DataFrame(data)
        for _, row in draw_df.iterrows():
            dwg.add(
                dwg.path(
                    d=row.path,
                    stroke=row.color,
                    fill="none",
                    stroke_width=row.width,
                    style="font-weight:bolder;",
                )
            )
        print(f"third layer: to_y2:{to_y2}")
        calc_node_arr = [
            {
                "name": to,
                "bottom": complex(to_x, to_y2 + total_width),
                "top": complex(to_x, to_y2),
            }
        ]
        self.calc_node_df = pd.concat([self.calc_node_df, pd.DataFrame(calc_node_arr)])

    def run_end_layer(self, to):
        dwg = self.dwg
        line_df = (
            self.line_df[(self.line_df.layer == 4) & (self.line_df.to == to)]
            .copy()
            .reset_index(drop=True)
        )
        line_df.val = line_df.val / RATIO
        calc_node_df = self.calc_node_df.copy()

        line_df = line_df.merge(calc_node_df, left_on="from", right_on="name")
        line_df = line_df.merge(self.node_df[["name", "color"]], on="name")
        node_df = self.end_node_df.copy()
        to_row = node_df[node_df.name == to].iloc[0]
        to_x, to_y = to_row.x, self.end_y

        data = []

        for i, row in line_df.iterrows():
            tmp_used_df = (
                self.used_df[self.used_df.name == row["from"]]
                .copy()
                .reset_index(drop=True)
            )
            start = row.top
            if tmp_used_df.shape[0] != 0:
                last_index = tmp_used_df.shape[0] - 1
                start = tmp_used_df.loc[last_index].bottom
            tmp_used_df = pd.DataFrame(
                [
                    {
                        "name": row["name"],
                        "top": start,
                        "bottom": start + complex(0, row.val),
                    }
                ]
            )
            self.used_df = pd.concat([self.used_df, tmp_used_df], ignore_index=True)

            p0 = start + complex(0, row.val / 2.0)
            p1 = complex(2800 + i * 10, start.imag + row.val / 2.0)
            p4 = complex(to_x, to_y + row.val / 2.0)
            print(f"p1:{p1}, p4:{p4}")
            p2 = complex(p1.real / 2.0 + p4.real / 2.0, p1.imag)
            p3 = complex(p1.real / 2.0 + p4.real / 2.0, p2.imag)
            p5 = p4 + complex(200, 0)
            seg0 = Line(p0, p1)
            seg1 = CubicBezier(p1, p2, p3, p4)
            seg2 = Line(p4, p5)
            path = Path(seg0, seg1, seg2)
            path = smoothed_path(path, tightness=0.5)
            to_y = to_y + row.val + MARGIN
            # spath = smoothed_path(path, tightness=0.5)
            print("\r\n")
            data.append(
                {
                    "path": path.d(),
                    "width": row.val,
                    "color": row.color,
                    "name": row["name"],
                }
            )
            # if i == 2:
            #     print(row)
            #     break
        draw_df = pd.DataFrame(data)
        for _, row in draw_df.iterrows():
            dwg.add(
                dwg.path(
                    d=row.path, stroke=row.color, fill="none", stroke_width=row.width
                )
            )

        seg0 = Line(
            complex(to_x + 230, self.end_y - MARGIN),
            complex(to_x + 240, self.end_y - MARGIN),
        )
        seg1 = Line(complex(to_x + 240, self.end_y - MARGIN), complex(to_x + 240, to_y))
        seg2 = Line(complex(to_x + 240, to_y), complex(to_x + 230, to_y))
        right_path = Path(seg0, seg1, seg2)
        dwg.add(dwg.path(d=right_path.d(), stroke="black", fill="none", stroke_width=1))
        dwg.add(
            dwg.text(
                to,
                insert=(to_x + 260, self.end_y / 2.0 + to_y / 2.0),
                style="font-weight:bolder;",
            )
        )
        self.end_y = to_y + MARGIN * 1.5

    def run(self):
        self.run_first_layer(self.df[self.df.cat == "天然气"], "天然气", "天然气汇合点")
        self.run_first_layer(self.df[self.df.cat == "一次电力"], "一次电力", "一次电力汇合点")
        self.run_first_layer(self.df[self.df.cat == "外来电"], "外来电", "外来电汇合点")
        self.run_first_layer(self.df[self.df.cat == "煤炭"], "煤炭", "煤炭汇合点")
        self.run_first_layer(self.df[self.df.cat == "石油"], "石油", "石油汇合点")
        self.run_second_layer("发电供热")
        print("-" * 100)
        self.run_second_layer("洗练")
        print("-" * 100)
        self.run_second_layer("炼化")
        print("-" * 100)
        self.run_third_layer("电力汇合点")
        print("end layer:" + "-" * 100)
        self.run_end_layer("工业")
        print("end layer:" + "-" * 100)
        self.run_end_layer("交通")
        print("end layer:" + "-" * 100)
        self.run_end_layer("建筑")
        print("end layer:" + "-" * 100)
        self.run_end_layer("其他")
        self.dwg.save(pretty=True)


engine = DrawEngine()
engine.run()
